'''
Author: SlytherinGe
LastEditTime: 2020-12-06 22:05:26
'''
# torch
import torch
import torch.utils.data as data
# utility
import numpy as np
import cv2
from tqdm import trange
from random import randint


class RamdonInstanceDataset(data.Dataset):

    def __init__(self, data_size, img_size, bg_color, target_color, random_noise=None, transform=None):

        self.data_size = data_size
        self.img_size = img_size
        self.bg_color = bg_color
        self.targeta_color = target_color
        self.random_noise = random_noise 
        self.transform = transform
        self.imgs = []
        self.masks = []
        print('start to generate pseudo instance targets')
        shortside_len = min(self.img_size)
        for _ in trange(self.data_size, colour='green'):
            img = np.ones((self.img_size[0], self.img_size[1], 3), dtype=np.uint8)
            img = img * np.array(self.bg_color, dtype=np.uint8)
            num_instance = 2#randint(1, 5)
            mask = np.zeros((num_instance, self.img_size[0], self.img_size[1]), dtype=np.uint8)
            instance_info = []
            for i in range(num_instance):
                valid_flag = False
                while not valid_flag:
                    valid_flag = True
                    pos = np.random.randint(0, shortside_len, size=2, dtype=np.int)
                    r = randint(int(shortside_len / num_instance * 0.49), int(shortside_len / num_instance * 0.5))
                    for ins_info in instance_info:
                        pos0, r0 = ins_info
                        dist = np.math.sqrt(sum((pos0-pos)**2))
                        if dist < r + r0:
                            valid_flag = False
                            break
                instance_info.append((pos, r))
                img = cv2.circle(img, center=(pos[0], pos[1]), radius=r, color=self.targeta_color, thickness=-1)
                mask[i] = cv2.circle(mask[i], center=(pos[0], pos[1]), radius=r, color=1, thickness=-1)
            if self.random_noise is not None:
                raise NotImplementedError
            self.imgs.append(img)
            self.masks.append(mask)


    def __len__(self):
        return self.data_size

    def __getitem__(self, index):
        
        img = self.imgs[index]
        mask = self.masks[index]

        if self.transform is not None:
            img = self.transform(img)

        return torch.from_numpy(img).permute(2,0,1).type(torch.float32), torch.from_numpy(mask).type(torch.float32)

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp))) 
        return fmt_str

if __name__ == '__main__':
    dataset = RamdonInstanceDataset(1024, (224, 224), (0,0,0), (255,255,255))